import os
import os.path as osp
import argparse
import onnx
from onnx_tf.backend import prepare


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--onnx_path', type=str, default=osp.join('data', 'qat.onnx'))
    parser.add_argument('--tf_path', type=str, default=osp.join('data', 'qat.pb'))
    args = parser.parse_args()

    onnx_model = onnx.load(args.onnx_path)
    tf_rep = prepare(onnx_model)
    tf_rep.export_graph(args.tf_path)


if __name__ == '__main__':
    main()
